import tensorflow as tf
from transformers import TFAutoModelForSeq2SeqLM

class MyOwnModel(tf.Module):
    def __init__(self, model_path="t5-small"):
        super(MyOwnModel, self).__init__()
        self.model = TFAutoModelForSeq2SeqLM.from_pretrained(model_path)

    @tf.function(input_signature=(tf.TensorSpec((None, 32), tf.int32, name="input_ids"),
                                  tf.TensorSpec((None, 32), tf.int32, name="attention_mask")), jit_compile=True)
    def serving(self, input_ids, attention_mask):
        outputs = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=32,
                                   return_dict_in_generate=True)
        return {"sequences": outputs["sequences"]}

model = MyOwnModel()
export_dir = "./"
tf.saved_model.save(
    model,
    export_dir,
    signatures={
        "serving_default":
            model.serving
    })

import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForSeq2SeqLM
export_dir = "./"
model = tf.saved_model.load(export_dir)

tokenizer = AutoTokenizer.from_pretrained("t5-small")
tokenization_kwargs = {"pad_to_multiple_of": 32, "padding": True, "return_tensors": "tf"}

input_prompts = [
    f"translate English to {language}: I have four cats and three dogs."
    for language in ["German", "French", "Romanian"]
]

def generate_text(inputs):
    tokenized_inputs = tokenizer(inputs, **tokenization_kwargs)
    generated_texts = model.signatures["serving_default"](**tokenized_inputs)
    for text in generated_texts["sequences"]:
        print(tokenizer.decode(text, skip_special_tokens=True))
# The first prompt will be slow (compiling), the others will be very fast!
generate_text(input_prompts[:2])
generate_text(input_prompts[:3])

import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("t5-small")
model = TFAutoModelForSeq2SeqLM.from_pretrained("t5-small")

# Main changes with respect to the original generate workflow: `tf.function` and `pad_to_multiple_of`
xla_generate = tf.function(model.generate, jit_compile=True)
tokenization_kwargs = {"pad_to_multiple_of": 32, "padding": True, "return_tensors": "tf"}

# The first prompt will be slow (compiling), the others will be very fast!
input_prompts = [
    f"translate English to {language}: I have four cats and three dogs."
    for language in ["German", "French", "Romanian"]
]
tokenized_inputs = tokenizer(input_prompts, **tokenization_kwargs)
generated_texts = xla_generate(**tokenized_inputs, max_new_tokens=32)
for text in generated_texts:
    print(tokenizer.decode(text, skip_special_tokens=True))

def serving(self, input_ids, attention_mask):
        outputs = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=32,
                                   return_dict_in_generate=True, num_beams=3, num_return_sequences=3)
        return {"sequences": outputs["sequences"]}